We define the following SIRS disease model and use a numerical solver to solve the differential equations for S, I and R:

# deSolve is used to numerically solve ODEs
library(deSolve)

# Define the SIRSS Disease model. Note x1 = alpha_SI, x2 = alpha_IR, x3 = alpha_SR .
# S = number of susceptible people, I = number of infected people, R = number of recovered people.
SIRS_Disease_Model <- function(t, y, parms) {       
  with(as.list(parms),{
    N  <-  y["S"] + y["I"] + y["R"]
    dS      =   x3 * y["R"]  -  x1 * y["S"] * y["I"] / N
    dI      =   x1 * y["S"] * y["I"] / N  - x2 * y["I"]
    dR    =   x2 * y["I"]  -  x3 * y["R"]
    res <- c(dS, dI, dR)
    list(res)
  })
}

We set some initial parameters (midrange values), an initial configuration and a time period up from t=0 to t=100:

# Set input parameters to their midrange values
parms = c(  x1 = 0.45, x2 = 0.25, x3 = 0.025 )

# Initial configuration for the number of individuals in compartment S, I and R at time t=0
ystart <- c(S = 850, I = 150, R = 0)                

# Define the time period
times <- seq(0, 100, length=1001)

Solve the differential equations and plot the output:

# Run the lsoda solver to solve the differential equations in the SIRS model
out <- as.matrix(lsoda(ystart, times, SIRS_Disease_Model, parms, maxsteps=20000))  

# "out" has first column time points, 2nd S, 3rd I and 4th R number
head(out)
##      time        S        I         R
## [1,]  0.0 850.0000 150.0000  0.000000
## [2,]  0.1 844.2488 151.9811  3.770086
## [3,]  0.2 838.4716 153.9484  7.580056
## [4,]  0.3 832.6700 155.9005 11.429445
## [5,]  0.4 826.8461 157.8361 15.317763
## [6,]  0.5 821.0017 159.7538 19.244481
# Plot results of the model output
plot_run <- function(){
  plot(out[,"time"],out[,"S"],lwd=6,col=4,ty="l",ylim=c(0,max(out[,-1])),xlab="time (t)",ylab="Number of People")
  lines(out[,"time"],out[,"I"],lwd=6,col=2)
  lines(out[,"time"],out[,"R"],lwd=6,col=3)
  legend("topright",legend=c("S = Susceptible"," I = Infected","R = Recovered"),col=c(4,2,3),lwd=6)
}

plot_run()

Bayes linear emulator without using any partial derivative information:

simple_BL_emulator_v2 <- function(x,              # the emulator prediction point
                                  xD,             # the run input locations xD
                                  D,              # the run outputs D = (f(x^1),...,f(x^n))
                                  theta = 1,      # the correlation lengths
                                  sigma = 1,      # the prior SD sigma sqrt(Var[f(x)])
                                  E_f = 0         # prior expectation of f: E(f(x)) = 0 
){
  
  # store length of runs D  
  n <- length(D)
  
  # Define Covariance structure of f(x): Cov[f(x),f(xdash)] 
  Cov_fx_fxdash <- function(x,xdash) sigma^2 * exp(-sum((x-xdash)^2)/theta^2) 
  
  
  # Define objects needed for BL update rules 
  # Create E[D] vector
  E_D <- rep(E_f,n)
  
  # Create Var_D matrix:
  Var_D <- matrix(0,nrow=n,ncol=n)

  for(i in 1:n) for(j in 1:n) Var_D[i,j] <- Cov_fx_fxdash(xD[i,],xD[j,])  
  
  # Create E[f(x)]
  E_fx <- E_f
  
  # Create Var_f(x) 
  Var_fx <- sigma^2
  
  # Create Cov_fx_D row vector
  Cov_fx_D <- matrix(0,nrow=1,ncol=n)

  for(j in 1:n) Cov_fx_D[1,j] <- Cov_fx_fxdash(x,xD[j,])    
  
  
  # Perform Bayes Linear adjustment to find Adjusted Expectation and Variance of f(x)
  ED_fx   <-  E_fx + Cov_fx_D %*% solve(Var_D) %*% (D - E_D)   
  VarD_fx <-  Var_fx - Cov_fx_D %*% solve(Var_D) %*% t(Cov_fx_D)  
  
  # Return the emulator adjusted expectation and variance 
  return(c("ExpD_f(x)"=ED_fx,"VarD_f(x)"=VarD_fx))  
  
}

Create a 16-point maximin Latin hypercube design:

lhd_maximin <- function(nl){                    # nl = number of points in LHD 
  
  x_lhd <- cbind("x1"=sample(0:(nl-1)),"x2"=sample(0:(nl-1))) / nl  +  0.5/nl  # create LHD
  
  ### Maximin loop: performs swaps on 1st of two closest points with another random point
  for(i in 1:1000){
    mat <- as.matrix(dist(x_lhd)) + diag(10,nl) # creates matrix of distances between points
    # note the inflated diagonal 
    closest_runs <- which(mat==min(mat),arr.ind=TRUE)   # finds pairs of closest runs
    ind <- closest_runs[sample(nrow(closest_runs),1),1] # chooses one of close runs at random
    swap_ind <- sample(setdiff(1:nl,ind),1)       # randomly selects another run to swap with
    x_lhd2 <- x_lhd                               # creates second version of LHD
    x_lhd2[ind[1],1]   <- x_lhd[swap_ind,1] # swaps x_1 vals between 1st close run & other run
    x_lhd2[swap_ind,1] <- x_lhd[ind[1],1]   # swaps x_1 vals between 1st close run & other run
    if(min(dist(x_lhd2)) >= min(dist(x_lhd))-0.00001) {  # if min distance between points is same or better
      x_lhd <- x_lhd2                                    # we replace LHD with new LHD with the swap
      # cat("min dist =",min(dist(x_lhd)),"Iteration = ",i,"\n") # write out min dist 
    }
  }
  
  
  return(x_lhd)
}

set.seed(15)
x_lhd<- lhd_maximin(16)
xD <- x_lhd

### plot maximin Latin hypercube design

plot(xD,xlim=c(0,1),ylim=c(0,1),pch=16,xaxs="i",yaxs="i",col="blue",xlab="x1",ylab="x2",cex=1.4)
abline(h=(0:16)/16,col="grey60")
abline(v=(0:16)/16,col="grey60")

We want to emulate over the 2-dimensional input space of \(x_1\) and \(x_2\) at time t=10, keeping \(x_3=0.04\). Here \(x_1 \in [0.1,0.8]\) and \(x_2 \in [0,0.5]\) so we scale these input ranges to [0,1] for our emulation.

xD_scaled <- cbind("x1"=rep(0,16),"x2"=rep(0,16))
xD_scaled[,1] <- xD[,1]*0.7+0.1
xD_scaled[,2] <- xD[,2]*0.5

We evaluate the true SIRS model for our 16 design runs and extract the model output at t=10 (this corresponds to the 101st row of the output matrix):

# Perform 16 runs of the SIRS model extracting the output for t=10 and store as D 
D <- NULL

for(i in 1:nrow(xD_scaled)){
  parms = c( xD_scaled[i,1], xD_scaled[i,2], x3 = 0.04)
  # Extract the output at t=10, this corresponds to the 101st row 
  infected <- as.matrix(lsoda(ystart, times, SIRS_Disease_Model, parms, maxsteps=20000)) [101,"I"]
  D <- c(D,infected)
  
}       

Define 50x50 grid of prediction points xP for input into our emulator:

x_grid <- seq(-0.001,1.001,len=50)
xP <- as.matrix(expand.grid("x1"=x_grid,"x2"=x_grid))

Emulate over this grid of prediction points and extract the expectation and variance:

# Evaluate emulator over 50x50=2500 prediction points xP
em_out <- t(apply(xP,1,simple_BL_emulator_v2,xD=xD,D=D,theta=0.25,sigma=250,E_f=350))   
E_D_fx_mat <- matrix(em_out[,"ExpD_f(x)"],nrow=length(x_grid),ncol=length(x_grid)) 
Var_D_fx_mat <- matrix(em_out[,"VarD_f(x)"],nrow=length(x_grid),ncol=length(x_grid)) 

Here is the plotting function for our output:

### define filled contour plot function for emulator output ###
emul_fill_cont <- function(
    cont_mat,            # matrix of values we want contour plot of 
    cont_levs=NULL,      # contour levels (NULL: automatic selection)
    nlev=20,             # approx no. of contour levels for auto select  
    plot_xD=TRUE,        # plot the design runs TRUE or FALSE
    xD=NULL,             # the design points if needed
    xD_col="green",      # colour of design runs
    x_grid,              # grid edge locations that define xP
    ...                  # extra arguments passed to filled.contour
){
  
  ### Define contour levels if necessary ###
  if(is.null(cont_levs)) cont_levs <- pretty(cont_mat,n=nlev)     
  
  ### create the filled contour plot ###
  filled.contour(x_grid,x_grid,cont_mat,levels=cont_levs,xlab="x1",ylab="x2",...,  
                 plot.axes={axis(1);axis(2)                 # sets up plotting in contour box
                   contour(x_grid,x_grid,cont_mat,add=TRUE,levels=cont_levs,lwd=0.8)   # plot contour lines
                   if(plot_xD) points(xD,pch=21,col=1,bg=xD_col,cex=1.5)})  # plot design points
}

Colour schemes:

library(viridisLite)

exp_cols <- plasma
var_cols <-  function(n) hcl.colors(n, "blues3", rev = TRUE)
diag_cols <- turbo

Plot the emulator adjusted expectation and variance:

emul_fill_cont(cont_mat=E_D_fx_mat,cont_levs=seq(-50,1000,50),xD=xD,x_grid=x_grid,
               color.palette=exp_cols,        # this sets the colour scheme
               main="Emulator Adjusted Expectation E_D[f(x)]")

emul_fill_cont(cont_mat=Var_D_fx_mat,cont_levs=NULL,xD=xD,x_grid=x_grid,
               color.palette=var_cols,
               main="Emulator Adjusted Variance Var_D[f(x)]")

In this instance we can plot the true 2-dimensional output of our SIRS model:

f <- NULL
for(i in 1:nrow(xP)){
  
  parms = c( xP[i,1]*0.7+0.1, xP[i,2]*0.5, x3 = 0.04)
  out <- as.matrix(lsoda(ystart, times, SIRS_Disease_Model, parms, maxsteps=20000)) [101,"I"]
  f <-  c(f,out)
}
fxP_mat <- matrix(f,nrow=length(x_grid),ncol=length(x_grid)) 



emul_fill_cont(cont_mat=fxP_mat,cont_levs=seq(-50,1000,50),xD=xD,x_grid=x_grid,
               color.palette=exp_cols,
               main="True Computer Model Function f(x)" )

Plotting the diagnostics:

S_diag_mat <- (E_D_fx_mat - fxP_mat) / sqrt(Var_D_fx_mat)

emul_fill_cont(cont_mat=S_diag_mat,cont_levs=seq(-3,3,0.25),xD=xD,x_grid=x_grid,
               xD_col="purple",
               color.palette=diag_cols,
               main="Emulator Diagnostics S_D[f(x)]")

We can see that the diagnostics look fine!

Derivative information can be extracted from the following adjoint SIRSS model:

### Define the SIRS Disease model. Note x1 = alpha_SI, x2 = alpha_IR, x3 = alpha_SR 
### S = number of susceptible people, I = number of infected people, R = number of recovered people.
SIRS_Disease_Model_adjoint <- function(t, y, parms) {   

  with(as.list(parms),{
    
 
    N  <-  y["S"] + y["I"] + y["R"]
    dS      =   x3 * y["R"]  -  x1 * y["S"] * y["I"] / N
    dI      =   x1 * y["S"] * y["I"] / N  - x2 * y["I"]
    dR    =   x2 * y["I"]  -  x3 * y["R"]
    
    dSdx1 = x3*y["dRdx1"] -x1*y["dSdx1"]*y["I"]/N - x1*y["S"]*y["dIdx1"]/N - y["S"]*y["I"]/N
    dSdx2 = x3*y["dRdx2"] -x1*y["dSdx2"]*y["I"]/N - x1*y["S"]*y["dIdx2"]/N
    dSdx3 = x3*y["dRdx3"]+ y["R"] -x1*y["dSdx3"]*y["I"] /N - x1*y["S"]*y["dIdx3"]/N
    
    dIdx1 = x1*y["dSdx1"]*y["I"]/N +x1*y["S"]*y["dIdx1"]/N +y["S"]*y["I"]/N -x2*y["dIdx1"]
    dIdx2 = x1*y["dSdx2"]*y["I"]/N + x1*y["S"]*y["dIdx2"]/N -x2*y["dIdx2"] - y["I"]
    dIdx3 = x1*y["dSdx3"]*y["I"] /N + x1*y["S"]*y["dIdx3"]/N - x2*y["dIdx3"]
    
    dRdx1 = x2*y["dIdx1"] - x3*y["dRdx1"]
    dRdx2 = y["I"] + x2*y["dIdx2"] -x3*y["dRdx2"]
    dRdx3 = x2*y["dIdx3"] - y["R"] -x3*y["dRdx3"]
 
    
    res <- c(dS, dI, dR, dSdx1,dSdx2,dSdx3,dIdx1,dIdx2,dIdx3,dRdx1,dRdx2,dRdx3)
    list(res)
  })
}

Bayes linear emulator using partial derivative information in both the \(x_1\) and \(x_2\) directions:

simple_BL_emulator_v2_dev<- function(x,              # the emulator prediction point
                                  xD,             # the run input locations xD
                                  D,              # the run outputs D = (f(x^1),...,f(x^n))
                                  theta = 1,      # the correlation lengths
                                  sigma = 1,      # the prior SD sigma sqrt(Var[f(x)])
                                  E_f=0,         # prior expectation of f: E(f(x)) = 0 
                                  n=16,           # # the number of design runs
                                  n_x1=16,       # the number of x1 derivatives 
                                  n_x2=16         # the number of x2 derivatives 
){
  


  # Define Covariance structure of f(x): Cov[f(x),f(xdash)] 
  Cov_fx_fxdash <- function(x,xdash) sigma^2 * exp(-sum((x-xdash)^2)/theta^2)
  
  # Derivatives here are w.r.t x1
  # Define Covariance structure of f(x): Cov[f'(x),f(xdash)] 
  Cov_fx_fxdash_dx1 <- function(x,xdash) -2*sigma^2 *(x[1]-xdash[1]) *exp(-sum((x-xdash)^2)/theta^2) /theta^2
  # Define Covariance structure of f(x): Cov[f(x),f'(xdash)] 
  Cov_fx_fxdash_dx1dash <- function(x,xdash) 2*sigma^2 *(x[1]-xdash[1]) *exp(-sum((x-xdash)^2)/theta^2)/theta^2
  # Define Covariance structure of f(x): Cov[f'(x),f'(xdash)] 
  Cov_fx_fxdash_dx1dx1dash <- function(x,xdash) -4*sigma^2 *(x[1]-xdash[1])^2 *exp(-sum((x-xdash)^2)/theta^2)/theta^4 + 2*sigma^2*exp(-sum((x-xdash)^2)/theta^2) /theta^2
  
  # Derivatives here are w.r.t x2
  # Define Covariance structure of f(x): Cov[f'(x),f(xdash)] 
  Cov_fx_fxdash_dx2 <- function(x,xdash) -2*sigma^2 *(x[2]-xdash[2]) *exp(-sum((x-xdash)^2)/theta^2) /theta^2
  # Define Covariance structure of f(x): Cov[f(x),f'(xdash)] 
  Cov_fx_fxdash_dx2dash <- function(x,xdash) 2*sigma^2 *(x[2]-xdash[2]) *exp(-sum((x-xdash)^2)/theta^2)/theta^2
  # Define Covariance structure of f(x): Cov[f'(x),f'(xdash)] 
  Cov_fx_fxdash_dx2dx2dash <- function(x,xdash) -4*sigma^2 *(x[2]-xdash[2])^2 *exp(-sum((x-xdash)^2)/theta^2)/theta^4 + 2*sigma^2*exp(-sum((x-xdash)^2)/theta^2) /theta^2
  
  # Mixed partial derivatives  
  Cov_mixed <- function(x,xdash) -4*sigma^2 *(x[2]-xdash[2])*(x[1]-xdash[1]) *exp(-sum((x-xdash)^2)/theta^2)/theta^4
  

  
  # Create E[D] vector
  # Give the derivatives zero expectation
  E_D <- c(rep(E_f,n),rep(0,n_x1),rep(0,n_x2))
  #E_D <- c(rep(E_f,n+n_x1+n_x2))
  
  # Create Var_D matrix
  Var_D <- matrix(0,nrow=n+n_x1+n_x2,ncol=n+n_x1+n_x2)
 
  # Keep this part of the matrix the same as in the non-derivative case
  for(i in 1:n) for(j in 1:n) Var_D[i,j] <- Cov_fx_fxdash(xD[i,],xD[j,]) 
  
  # Include the derivatives w.r.t x1
  for(i in 1:n) for(j in (n+1):(n+n_x1)) Var_D[i,j] <- Cov_fx_fxdash_dx1dash(xD[i,],xD[j,])  
  
  for(i in (n+1):(n+n_x1)) for(j in 1:n) Var_D[i,j] <-Cov_fx_fxdash_dx1(xD[i,],xD[j,])  
  
  for(i in (n+1):(n+n_x1)) for(j in (n+1):(n+n_x1) ) Var_D[i,j] <-Cov_fx_fxdash_dx1dx1dash(xD[i,],xD[j,]) 
  
  
  # Now include the derivatives w.r.t x2
  for(i in 1:n) for(j in (n+n_x1+1):(n+n_x1+n_x2)) Var_D[i,j] <- Cov_fx_fxdash_dx2dash(xD[i,],xD[j,])  
  
  for(i in (n+n_x1+1):(n+n_x1+n_x2)) for(j in 1:n)  Var_D[i,j] <-Cov_fx_fxdash_dx2(xD[i,],xD[j,])  
  
  for(i in (n+n_x1+1):(n+n_x1+n_x2)) for(j in (n+n_x1+1):(n+n_x1+n_x2) ) Var_D[i,j] <- Cov_fx_fxdash_dx2dx2dash(xD[i,],xD[j,])
  
  for(i in (n+n_x1+1):(n+n_x1+n_x2)) for(j in (n+1):(n+n_x1)) Var_D[i,j] <- Cov_mixed(xD[i,],xD[j,])
    
  for(i in (n+1):(n+n_x1))   for(j in (n+n_x1+1):(n+n_x1+n_x2)) Var_D[i,j] <- Cov_mixed(xD[i,],xD[j,])

  # Create E[f(x)]
  E_fx <- E_f
  
  # Create Var_f(x) 
  Var_fx <- sigma^2
  
  # Create Cov_fx_D row vector
  Cov_fx_D <- matrix(0,nrow=1,ncol=n+n_x1+n_x2)
  
  # Covariance for our known runs
  for(j in 1:n) Cov_fx_D[1,j] <- Cov_fx_fxdash(x,xD[j,])    
  # Covariance for x1 partial derivatives
  for(j in (n+1):(n+n_x1)) Cov_fx_D[1,j] <- Cov_fx_fxdash_dx1dash(x,xD[j,])
  # Covariance for x2 partial derivatives
  for(j in (n+n_x1+1):(n+n_x1+n_x2)) Cov_fx_D[1,j] <- Cov_fx_fxdash_dx2dash(x,xD[j,])
  

  #Perform Bayes Linear adjustment to find Adjusted Expectation and Variance of f(x)
  ED_fx   <-  E_fx + Cov_fx_D %*% solve(Var_D) %*% (D - E_D)   
  VarD_fx <-  Var_fx - Cov_fx_D %*% solve(Var_D) %*% t(Cov_fx_D)  
  
  ### Return the emulator adjusted expectation and variance

  return(c("ExpD_f(x)"=ED_fx,"VarD_f(x)"=VarD_fx))  

}

Calculate all partial derivatives in the \(x_1\) and \(x_2\) direction:

# Starting configuration for the SIRSS model 

ystart <- c(S = 850, I = 150, R = 0, dSdx1=0, dSdx2=0, dSdx3=0, dIdx1=0, dIdx2=0, dIdx3=0, dRdx1=0, dRdx2=0, dRdx3=0)


x1_dev <- NULL
for(i in 1:nrow(xD)){
  parms = c(xD_scaled[i,1], xD_scaled[i,2], x3=0.04,dSdx1 = 0, dIdx1 = 0, dRdx1 = 0, dSdx2 = 0, dIdx2 = 0, dRdx2 = 0, dSdx3 = 0, dIdx3 = 0, dRdx3 = 0)
  # Extract the output at t=10, this corresponds to the 101st row 
  infected <- as.matrix(lsoda(ystart, times, SIRS_Disease_Model_adjoint, parms, maxsteps=200000)) [101,"dIdx1"]
  x1_dev <- c(x1_dev,infected)
  
}       

x1_dev <- x1_dev*0.7
# To get original scale just multiply by 0.7

x2_dev <- NULL
for(i in 1:nrow(xD_scaled)){
  parms = c(xD_scaled[i,1], xD_scaled[i,2], x3=0.04, dSdx1 = 0, dIdx1 = 0, dRdx1 = 0, dSdx2 = 0, dIdx2 = 0, dRdx2 = 0, dSdx3 = 0, dIdx3 = 0, dRdx3 = 0)
  # Extract the output at t=10, this corresponds to the 101st row 
  infected <- as.matrix(lsoda(ystart, times, SIRS_Disease_Model_adjoint, parms, maxsteps=20000)) [101,"dIdx2"]
  x2_dev <- c(x2_dev,infected)
  
}   

x2_dev <- x2_dev*0.5

Emulate using this complete derivative information:

# Add this derivative information to D
D <- c(D,x1_dev,x2_dev)
# Record the input points for each partial derivative
xD <- rbind(xD,xD,xD)

em_out <- t(apply(xP,1,simple_BL_emulator_v2_dev,xD=xD,D=D,theta=0.25,sigma=250,E_f=350))


### store emulator output as matrices to aid plotting ###
E_D_fx_mat <- matrix(em_out[,"ExpD_f(x)"],nrow=length(x_grid),ncol=length(x_grid)) 
Var_D_fx_mat <- matrix(em_out[,"VarD_f(x)"],nrow=length(x_grid),ncol=length(x_grid)) 

emul_fill_cont(cont_mat=E_D_fx_mat,cont_levs=seq(-50,1000,50),xD=xD,x_grid=x_grid,
               color.palette=exp_cols,        # this sets the colour scheme
               main="Emulator Adjusted Expectation E_D[f(x)]")

emul_fill_cont(cont_mat=Var_D_fx_mat,cont_levs=NULL,xD=xD,x_grid=x_grid,
               color.palette=var_cols,
               main="Emulator Adjusted Variance Var_D[f(x)]")

Plotting the diagnostics:

S_diag_mat <- (E_D_fx_mat - fxP_mat) / sqrt(Var_D_fx_mat)

emul_fill_cont(cont_mat=S_diag_mat,cont_levs=seq(-3,3,0.25),xD=xD,x_grid=x_grid,
               xD_col="purple",
               color.palette=diag_cols,
               main="Emulator Diagnostics S_D[f(x)]")

Considering only the derivative information in the \(x_1\) direction:

simple_BL_emulator_v2_x1 <- function(x,              # the emulator prediction point
                                  xD,             # the run input locations xD
                                  D,              # the run outputs D = (f(x^1),...,f(x^n))
                                  theta = 1,      # the correlation lengths
                                  sigma = 1,      # the prior SD sigma sqrt(Var[f(x)])
                                  E_f = 0,       # prior expectation of f: E(f(x)) = 0 
                                  n_x1 = 16   # the number of x1 derivatives 
){
  
  # store length of runs D  
  n <- 16
  
  
  ### Define Covariance structure of f(x): Cov[f(x),f(xdash)] 
  Cov_fx_fxdash <- function(x,xdash) sigma^2 * exp(-sum((x-xdash)^2)/theta^2)
  
  # Derivatives are w.r.t x1:
  # Define Covariance structure of f(x): Cov[f'(x),f(xdash)] 
  Cov_fx_fxdash_dx1 <- function(x,xdash) -2*sigma^2 *(x[1]-xdash[1]) *exp(-sum((x-xdash)^2)/theta^2) /theta^2
  # Define Covariance structure of f(x): Cov[f(x),f'(xdash)] 
  Cov_fx_fxdash_dx1dash <- function(x,xdash) 2*sigma^2 *(x[1]-xdash[1]) *exp(-sum((x-xdash)^2)/theta^2)/theta^2
  # Define Covariance structure of f(x): Cov[f'(x),f'(xdash)] 
  Cov_fx_fxdash_dx1dx1dash <- function(x,xdash) -4*sigma^2 *(x[1]-xdash[1])^2 *exp(-sum((x-xdash)^2)/theta^2)/theta^4 + 2*sigma^2*exp(-sum((x-xdash)^2)/theta^2) /theta^2
  
  
 
  # Create E[D] vector
  # Give the x1 partial derivatives zero expectation
  E_D <- c(rep(E_f,n), rep(0,n_x1))
  
  # Create Var_D matrix:
  Var_D <- matrix(0,nrow=n+n_x1,ncol=n+n_x1)
  
  # Keep this part of the matrix the same as emulating without derivative information
  for(i in 1:n) for(j in 1:n) Var_D[i,j] <- Cov_fx_fxdash(xD[i,],xD[j,])  
  
  # Including the x1 partial derivatives
  for(i in 1:n) for(j in (n+1):(n+n_x1)) Var_D[i,j] <- Cov_fx_fxdash_dx1dash(xD[i,],xD[j,])  
  
  for(j in 1:n) for(i in (n+1):(n+n_x1)) Var_D[i,j] <-Cov_fx_fxdash_dx1(xD[i,],xD[j,])  
  
  for(i in (n+1):(n+n_x1)) for(j in (n+1):(n+n_x1) )    Var_D[i,j] <-Cov_fx_fxdash_dx1dx1dash(xD[i,],xD[j,]) 
  
  
  
  # Create E[f(x)]
  E_fx <- E_f
  
  # Create Var_f(x) 
  Var_fx <- sigma^2
  
  # Create Cov_fx_D row vector
  Cov_fx_D <- matrix(0,nrow=1,ncol=n+n_x1)
  
  for(j in 1:n) Cov_fx_D[1,j] <- Cov_fx_fxdash(x,xD[j,])    
  
  # Include the x1 partial derivative information
  for(j in (n+1):(n+n_x1)) Cov_fx_D[1,j] <- Cov_fx_fxdash_dx1dash(x,xD[j,])

  
  #Perform Bayes Linear adjustment to find Adjusted Expectation and Variance of f(x)
  ED_fx   <-  E_fx + Cov_fx_D %*% solve(Var_D) %*% (D - E_D)   
  VarD_fx <-  Var_fx - Cov_fx_D %*% solve(Var_D) %*% t(Cov_fx_D)  

  # Return emulator adjusted expectation and variance 
  return(c("ExpD_f(x)"=ED_fx,"VarD_f(x)"=VarD_fx))  
  
  
}

Emulate only considering the partial derivatives for all of the known runs in the \(x_1\) direction:

ystart <- c(S = 850, I = 150, R = 0)


# Perform 16 runs of the SIR model extracting the output for t=10 and store as D 
D <- NULL

for(i in 1:nrow(xD_scaled)){
  parms = c( xD_scaled[i,1], xD_scaled[i,2], x3 = 0.04)
  infected <- as.matrix(lsoda(ystart, times, SIRS_Disease_Model, parms, maxsteps=20000)) [101,"I"]
  D <- c(D,infected)
  
}       


### Define 50x50 grid of prediction points xP for emulator evaluation ###
x_grid <- seq(-0.001,1.001,len=50)
xP <- as.matrix(expand.grid("x1"=x_grid,"x2"=x_grid))



D <- c(D,x1_dev)
xD <- rbind(xD,xD)
em_out <- t(apply(xP,1,simple_BL_emulator_v2_x1,xD=xD,D=D,n_x1=16,theta=0.25,sigma=250,E_f=350))


### store emulator output as matrices to aid plotting ###
E_D_fx_mat <- matrix(em_out[,"ExpD_f(x)"],nrow=length(x_grid),ncol=length(x_grid)) 
Var_D_fx_mat <- matrix(em_out[,"VarD_f(x)"],nrow=length(x_grid),ncol=length(x_grid)) 

emul_fill_cont(cont_mat=E_D_fx_mat,cont_levs=seq(-50,1000,50),xD=xD,x_grid=x_grid,
               color.palette=exp_cols,        # this sets the colour scheme
               main="Emulator Adjusted Expectation E_D[f(x)]")

emul_fill_cont(cont_mat=Var_D_fx_mat,cont_levs=NULL,xD=xD,x_grid=x_grid,
               color.palette=var_cols,
               main="Emulator Adjusted Variance Var_D[f(x)]")

Now just including all of the partial derivatives in the \(x_2\) direction:

simple_BL_emulator_v2_x2 <- function(x,              # the emulator prediction point
                                  xD,             # the run input locations xD
                                  D,              # the run outputs D = (f(x^1),...,f(x^n))
                                  theta = 1,      # the correlation lengths
                                  sigma = 1,      # the prior SD sigma sqrt(Var[f(x)])
                                  E_f = 0 ,      # prior expectation of f: E(f(x)) = 0 
                                  n_x2 = 16    # the number of x2 derivatives 
){
  
  # store length of runs D  
  n <- 16
  
  # Define Covariance structure of f(x): Cov[f(x),f(xdash)] 
  Cov_fx_fxdash <- function(x,xdash) sigma^2 * exp(-sum((x-xdash)^2)/theta^2)
 
  
  # Derivatives are w.r.t x2:
  # Define Covariance structure of f(x): Cov[f'(x),f(xdash)] 
  Cov_fx_fxdash_dx2 <- function(x,xdash) -2*sigma^2 *(x[2]-xdash[2]) *exp(-sum((x-xdash)^2)/theta^2) /theta^2
  # Define Covariance structure of f(x): Cov[f(x),f'(xdash)] 
  Cov_fx_fxdash_dx2dash <- function(x,xdash) 2*sigma^2 *(x[2]-xdash[2]) *exp(-sum((x-xdash)^2)/theta^2)/theta^2
  # Define Covariance structure of f(x): Cov[f'(x),f'(xdash)] 
  Cov_fx_fxdash_dx2dx2dash <- function(x,xdash) -4*sigma^2 *(x[2]-xdash[2])^2 *exp(-sum((x-xdash)^2)/theta^2)/theta^4 + 2*sigma^2*exp(-sum((x-xdash)^2)/theta^2) /theta^2
  
  
  # Create E[D] vector
  # Give the x2 partial derivatives zero expectation
  E_D <- c(rep(E_f,n),rep(0,n_x2))
  
  # Create Var_D matrix:
  Var_D <- matrix(0,nrow=n+n_x2,ncol=n+n_x2)
  
  # Keep this part of the matrix the same as emulating without derivative information
  for(i in 1:n) for(j in 1:n) Var_D[i,j] <- Cov_fx_fxdash(xD[i,],xD[j,])  
  
  # Now include the x2 partial derivatives 
  for(i in 1:n) for(j in (n+1):(n+n_x2)) Var_D[i,j] <- Cov_fx_fxdash_dx2dash(xD[i,],xD[j,])  
  
  for(j in 1:n) for(i in (n+1):(n+n_x2)) Var_D[i,j] <-Cov_fx_fxdash_dx2(xD[i,],xD[j,])  
  
  for(i in (n+1):(n+n_x2)) for(j in (n+1):(n+n_x2) )    Var_D[i,j] <-Cov_fx_fxdash_dx2dx2dash(xD[i,],xD[j,]) 
  
  
  # Create E[f(x)]
  E_fx <- E_f
  
  # Create Var_f(x) 
  Var_fx <- sigma^2
  
  # Create Cov_fx_D row vector
  Cov_fx_D <- matrix(0,nrow=1,ncol=n+n_x2)
  
  for(j in 1:n) Cov_fx_D[1,j] <- Cov_fx_fxdash(x,xD[j,])    
  
  for(j in (n+1):(n+n_x2)) Cov_fx_D[1,j] <- Cov_fx_fxdash_dx2dash(x,xD[j,])

  
  # Perform Bayes Linear adjustment to find Adjusted Expectation and Variance of f(x) 
  ED_fx   <-  E_fx + Cov_fx_D %*% solve(Var_D) %*% (D - E_D)   
  VarD_fx <-  Var_fx - Cov_fx_D %*% solve(Var_D) %*% t(Cov_fx_D)  

  # Return emulator adjusted expectation and variance 
  return(c("ExpD_f(x)"=ED_fx,"VarD_f(x)"=VarD_fx))  
  
}

Rebuilding the emulator:

xD <- x_lhd
# Perform 16 runs of the SIR model extracting the output for t=10 and store as D 
D <- NULL

for(i in 1:nrow(xD_scaled)){
  parms = c( xD_scaled[i,1], xD_scaled[i,2], x3 = 0.04)
  infected <- as.matrix(lsoda(ystart, times, SIRS_Disease_Model, parms, maxsteps=20000)) [101,"I"]
  D <- c(D,infected)
  
}       

# Define 50x50 grid of prediction points xP for emulator evaluation 
x_grid <- seq(-0.001,1.001,len=50)
xP <- as.matrix(expand.grid("x1"=x_grid,"x2"=x_grid))


D <- c(D,x2_dev)
xD <- rbind(xD,xD)
em_out <- t(apply(xP,1,simple_BL_emulator_v2_x2,xD=xD,D=D,theta=0.25,sigma=250,E_f=350))

# Store the emulator output as matrices 
E_D_fx_mat <- matrix(em_out[,"ExpD_f(x)"],nrow=length(x_grid),ncol=length(x_grid)) 
Var_D_fx_mat <- matrix(em_out[,"VarD_f(x)"],nrow=length(x_grid),ncol=length(x_grid)) 

emul_fill_cont(cont_mat=E_D_fx_mat,cont_levs=seq(-50,1000,50),xD=xD,x_grid=x_grid,
               color.palette=exp_cols,        # this sets the colour scheme
               main="Emulator Adjusted Expectation E_D[f(x)]")

emul_fill_cont(cont_mat=Var_D_fx_mat,cont_levs=NULL,xD=xD,x_grid=x_grid,
               color.palette=var_cols,
               main="Emulator Adjusted Variance Var_D[f(x)]")

Now emulate over the 2-dimensional input space of \(x_1\) and \(x_3\) at time t=10, keeping \(x_2=0.7\). Here \(x_1 \in [0.1,0.8]\) and \(x_3 \in [0,0.05]\) so we scale these input ranges to [0,1] for our emulation.

set.seed(29)
x_lhd <- lhd_maximin(8)
# Define run locations as the Maximin LHD 
xD <- x_lhd 

xD_scaled <- cbind("x2"=rep(0,8),"x3"=rep(0,8))
xD_scaled[,1] <- xD[,1]*0.5
xD_scaled[,2] <- xD[,2]*0.05

# Perform 16 runs of the SIR model extracting the output for t=10 and store as D 
D <- NULL

for(i in 1:nrow(xD_scaled)){
  parms = c( x1=0.7, xD_scaled[i,1], xD_scaled[i,2])
  # Extract the output at t=10, this corresponds to the 101st row 
  infected <- as.matrix(lsoda(ystart, times, SIRS_Disease_Model, parms, maxsteps=20000)) [101,"R"]
  D <- c(D,infected)
  
}       

Emulator which includes derivative information and different correlation length in each direction:

simple_BL_emulator_dev_theta<- function(x,              # the emulator prediction point
                                         xD,             # the run input locations xD
                                         D,              # the run outputs D = (f(x^1),...,f(x^n))
                                         theta = c(1,1),      # the correlation lengths
                                         sigma = 1,      # the prior SD sigma sqrt(Var[f(x)])
                                         E_f=0,         # prior expectation of f: E(f(x)) = 0 
                                         n=8,           # # the number of design runs
                                         n_x1=8,       # the number of x1 derivatives 
                                         n_x2=8         # the number of x2 derivatives 
){
  
  
  
  
  ### Define Covariance structure of f(x): Cov[f(x),f(xdash)] ###
  Cov_fx_fxdash <- function(x,xdash) sigma^2 * exp(-(x[1]-xdash[1])^2/theta[1]^2 -(x[2]-xdash[2])^2/theta[2]^2)
  
  ## Derivatives w.r.t x1 ##
  ### Define Covariance structure of f(x): Cov[f'(x),f(xdash)] ###
  Cov_fx_fxdash_dx1 <- function(x,xdash) -2*sigma^2 *(x[1]-xdash[1]) *exp(-(x[1]-xdash[1])^2/theta[1]^2 -(x[2]-xdash[2])^2/theta[2]^2) /theta[1]^2
  ### Define Covariance structure of f(x): Cov[f(x),f'(xdash)] ###
  Cov_fx_fxdash_dx1dash <- function(x,xdash) 2*sigma^2 *(x[1]-xdash[1]) *exp(-(x[1]-xdash[1])^2/theta[1]^2 -(x[2]-xdash[2])^2/theta[2]^2)/theta[1]^2
  ### Define Covariance structure of f(x): Cov[f'(x),f'(xdash)] ###
  Cov_fx_fxdash_dx1dx1dash <- function(x,xdash) -4*sigma^2 *(x[1]-xdash[1])^2 *exp(-(x[1]-xdash[1])^2/theta[1]^2 -(x[2]-xdash[2])^2/theta[2]^2)/theta[1]^4 + 2*sigma^2*exp(-(x[1]-xdash[1])^2/theta[1]^2 -(x[2]-xdash[2])^2/theta[2]^2) /theta[1]^2
  
  
  ## Derivatives w.r.t x2 ##
  ### Define Covariance structure of f(x): Cov[f'(x),f(xdash)] ###
  Cov_fx_fxdash_dx2 <- function(x,xdash) -2*sigma^2 *(x[2]-xdash[2]) *exp(-(x[1]-xdash[1])^2/theta[1]^2 -(x[2]-xdash[2])^2/theta[2]^2) /theta[2]^2
  ### Define Covariance structure of f(x): Cov[f(x),f'(xdash)] ###
  Cov_fx_fxdash_dx2dash <- function(x,xdash) 2*sigma^2 *(x[2]-xdash[2]) *exp(-(x[1]-xdash[1])^2/theta[1]^2 -(x[2]-xdash[2])^2/theta[2]^2)/theta[2]^2
  ### Define Covariance structure of f(x): Cov[f'(x),f'(xdash)] ###
  Cov_fx_fxdash_dx2dx2dash <- function(x,xdash) -4*sigma^2 *(x[2]-xdash[2])^2 *exp(-(x[1]-xdash[1])^2/theta[1]^2 -(x[2]-xdash[2])^2/theta[2]^2)/theta[2]^4 + 2*sigma^2*exp(-(x[1]-xdash[1])^2/theta[1]^2 -(x[2]-xdash[2])^2/theta[2]^2) /theta[2]^2
  
  # Mixed partial derivatives 
  Cov_mixed <- function(x,xdash) -4*sigma^2 *(x[2]-xdash[2])*(x[1]-xdash[1]) *exp(-(x[1]-xdash[1])^2/theta[1]^2 -(x[2]-xdash[2])^2/theta[2]^2)/(theta[1]^2*theta[2]^2)
  
  
  
  # Create E[D] vector
  E_D <- c(rep(E_f,n),rep(0,n_x1),rep(0,n_x2))
  #E_D <- c(rep(E_f,n+n_x1+n_x2))
  
  # Create Var_D matrix:
  Var_D <- matrix(0,nrow=n+n_x1+n_x2,ncol=n+n_x1+n_x2)
  
  # Keep this part of the matrix the same 
  for(i in 1:n) for(j in 1:n) Var_D[i,j] <- Cov_fx_fxdash(xD[i,],xD[j,]) 
  
  # Include the derivatives w.r.t x1
  for(i in 1:n) for(j in (n+1):(n+n_x1)) Var_D[i,j] <- Cov_fx_fxdash_dx1dash(xD[i,],xD[j,])  
  
  for(i in (n+1):(n+n_x1)) for(j in 1:n) Var_D[i,j] <-Cov_fx_fxdash_dx1(xD[i,],xD[j,])  
  
  for(i in (n+1):(n+n_x1)) for(j in (n+1):(n+n_x1) ) Var_D[i,j] <-Cov_fx_fxdash_dx1dx1dash(xD[i,],xD[j,]) 
  
  
  # Now include the derivatives w.r.t x2
  for(i in 1:n) for(j in (n+n_x1+1):(n+n_x1+n_x2)) Var_D[i,j] <- Cov_fx_fxdash_dx2dash(xD[i,],xD[j,])  
  
  for(i in (n+n_x1+1):(n+n_x1+n_x2)) for(j in 1:n)  Var_D[i,j] <-Cov_fx_fxdash_dx2(xD[i,],xD[j,])  
  
  for(i in (n+n_x1+1):(n+n_x1+n_x2)) for(j in (n+n_x1+1):(n+n_x1+n_x2) ) Var_D[i,j] <- Cov_fx_fxdash_dx2dx2dash(xD[i,],xD[j,])
  
  for(i in (n+n_x1+1):(n+n_x1+n_x2)) for(j in (n+1):(n+n_x1)) Var_D[i,j] <- Cov_mixed(xD[i,],xD[j,])
  
  for(i in (n+1):(n+n_x1))   for(j in (n+n_x1+1):(n+n_x1+n_x2)) Var_D[i,j] <- Cov_mixed(xD[i,],xD[j,])
  
  # Create E[f(x)]
  E_fx <- E_f
  
  # Create Var_f(x) 
  Var_fx <- sigma^2
  
  # Create Cov_fx_D row vector
  Cov_fx_D <- matrix(0,nrow=1,ncol=n+n_x1+n_x2)
  
  for(j in 1:n) Cov_fx_D[1,j] <- Cov_fx_fxdash(x,xD[j,])    
  
  for(j in (n+1):(n+n_x1)) Cov_fx_D[1,j] <- Cov_fx_fxdash_dx1dash(x,xD[j,])
  
  for(j in (n+n_x1+1):(n+n_x1+n_x2)) Cov_fx_D[1,j] <- Cov_fx_fxdash_dx2dash(x,xD[j,])
  
  
  ### Perform Bayes Linear adjustment to find Adjusted Expectation and Variance of f(x) ###
  ED_fx   <-  E_fx + Cov_fx_D %*% solve(Var_D) %*% (D - E_D)   
  VarD_fx <-  Var_fx - Cov_fx_D %*% solve(Var_D) %*% t(Cov_fx_D)  
  
  ### return emulator expectation and variance ###
  
  return(c("ExpD_f(x)"=ED_fx,"VarD_f(x)"=VarD_fx))  
  
}

Extract the derivative information from the adjoint model:

ystart <- c(S = 850, I = 150, R = 0, dSdx1=0, dSdx2=0, dSdx3=0, dIdx1=0, dIdx2=0, dIdx3=0, dRdx1=0, dRdx2=0, dRdx3=0)


x2_dev <- NULL
for(i in 1:nrow(xD)){
  parms = c(x1=0.7,xD_scaled[i,1], xD_scaled[i,2],dSdx1 = 0, dIdx1 = 0, dRdx1 = 0, dSdx2 = 0, dIdx2 = 0, dRdx2 = 0, dSdx3 = 0, dIdx3 = 0, dRdx3 = 0)
  # Extract the output at t=10, this corresponds to the 101st row 
  infected <- as.matrix(lsoda(ystart, times, SIRS_Disease_Model_adjoint, parms, maxsteps=200000)) [101,"dRdx2"]
  x2_dev <- c(x2_dev,infected)
  
}       

x2_dev <- x2_dev*0.5
# To get original scale just multiply by 0.5

x3_dev <- NULL
for(i in 1:nrow(xD)){
  parms = c(x1=0.7,xD_scaled[i,1], xD_scaled[i,2],dSdx1 = 0, dIdx1 = 0, dRdx1 = 0, dSdx2 = 0, dIdx2 = 0, dRdx2 = 0, dSdx3 = 0, dIdx3 = 0, dRdx3 = 0)
  # Extract the output at t=10, this corresponds to the 101st row 
  infected <- as.matrix(lsoda(ystart, times, SIRS_Disease_Model_adjoint, parms, maxsteps=200000)) [101,"dRdx3"]
  x3_dev <- c(x3_dev,infected)
  
}       

x3_dev <- x3_dev*0.05
# To get original scale just multiply by 0.5

Add in this derivative information and emulate using different correlation lenghts:

D <- c(D,x2_dev,x3_dev)
xD <- rbind(xD,xD,xD)
em_out <- t(apply(xP,1,simple_BL_emulator_dev_theta,xD=xD,D=D,theta=c(0.25,0.75),sigma=250,E_f=350))

E_D_fx_mat <- matrix(em_out[,"ExpD_f(x)"],nrow=length(x_grid),ncol=length(x_grid)) 
Var_D_fx_mat <- matrix(em_out[,"VarD_f(x)"],nrow=length(x_grid),ncol=length(x_grid)) 


emul_fill_cont(cont_mat=E_D_fx_mat,cont_levs=seq(0,700,50),xD=xD,x_grid=x_grid,
                color.palette=exp_cols,        # this sets the colour scheme
                main="Emulator Adjusted Expectation E_D[f(x)]")

emul_fill_cont(cont_mat=Var_D_fx_mat,cont_levs=NULL,xD=xD,x_grid=x_grid,
               color.palette=var_cols,
               main="Emulator Adjusted Variance Var_D[f(x)]")